#include "ComputeCalibrationRect.hpp"
#include "MiscFunctions.h"

#include <stdio.h>
#include <opencv2/calib3d/calib3d.hpp>
#include <algorithm>
#include <stdexcept>

using namespace cv;
using namespace std;
using namespace idch;


ComputeCalibrationRect::
        ComputeCalibrationRect( double targetWidth , double targetHeight ,
                                int imgWidth , int imgHeight , int numTrials, double inlierTol )
                                    : calib(4)
{
    this->imgWidth = imgWidth;
    this->imgHeight = imgHeight;
    this->inlierTol = inlierTol;
    this->numTrials = numTrials;
    this->errorUsed = 0;
    this->sampleSize = 15;

    //tl
    targetDesc.push_back(Point3f(0.0,targetHeight,0.0));
    //tr
    targetDesc.push_back(Point3f(targetWidth,targetHeight,0.0));
    //br
    targetDesc.push_back(Point3f(targetWidth,0.0,0.0));
    //bl
    targetDesc.push_back(Point3f(0.0,0.0,0.0));
}

ComputeCalibrationRect::~ComputeCalibrationRect() {

}

void ComputeCalibrationRect::
        process( const std::vector<CalibrationLocation>& obs ) {

//    std::vector<int> indexes;
//    for( size_t i = 0; i < obs.size(); ++i ) {
//        indexes.push_back(i);
//    }

//    std::vector<CalibrationLocation> smallSet;

//    int bestSize = 0;
//    StereoCalib bestModel( calib.numCoefs() );
//    StereoCalib model(calib.numCoefs());

//    for( int i = 0; i < numTrials; i++ ) {
//        // randomly sample a small number of points
//        std::random_shuffle( indexes.begin(), indexes.end() );

//        smallSet.clear();
//        for( size_t j = 0; j < sampleSize; ++j ) {
//            smallSet.push_back( obs.at(indexes.at(j)));
//        }

//        double error = calibrate(smallSet,model,false);

//        int inlierSize = countInliers(obs,model);

//        if( inlierSize > bestSize ) {
//            bestSize = inlierSize;
//            bestModel = model;
//        }

//        printf("%d Error = %f  inlier size = %d out of %d\n",i,error,inlierSize,(int)obs.size());

//    }

//    if( bestSize < sampleSize ) {
//        throw std::runtime_error("best set is too small.");
//    }

//    // create a list of all the inliers for the best model
//    std::vector<CalibrationLocation> bestInliers;
//    findInliers(obs,bestModel,bestInliers);

    // recompute the calibration using all the points
//    errorUsed = calibrate(bestInliers,bestModel,true);
    errorUsed = calibrate(obs,calib,false);

//    printf("NUmber of inliers %d\n",countInliers(obs,calib));

    printf("Projection error = %f\n",errorUsed);

    printf("Refining\n");
    std::vector<CalibrationLocation> best;
    pruneWorst(obs,calib,best,0.8);
    errorUsed = calibrate(best,calib,false);
    printf("Number of inliers %d\n",countInliers(obs,calib));
    printf("Projection error = %f\n",errorUsed);
}

double ComputeCalibrationRect::
        getError() {
    return errorUsed;
}

const StereoCalib& ComputeCalibrationRect::
        getCalib() {
    return calib;
}

const std::vector<CalibrationLocation>& ComputeCalibrationRect::
        getUsed() {
    return used;
}

double ComputeCalibrationRect::
        calibrate( const std::vector<CalibrationLocation>& obs ,
                   StereoCalib &calib , bool useEstimate) {
    // create the list of targets
    std::vector<std::vector<cv::Point3f> > objectPoints;

    for( size_t i = 0; i < obs.size(); ++i ) {
        objectPoints.push_back(targetDesc);
    }

    // list of observations
    vector<vector<Point2f> > depthObs;
    vector<vector<Point2f> > rgbObs;

    for( size_t i = 0; i < obs.size(); ++i ) {
        const CalibrationLocation& c = obs.at(i);

        vector<Point2f> d;
        vector<Point2f> r;

        d.push_back(c.depth.tl);
        d.push_back(c.depth.tr);
        d.push_back(c.depth.br);
        d.push_back(c.depth.bl);

        r.push_back(c.rgb.tl);
        r.push_back(c.rgb.tr);
        r.push_back(c.rgb.br);
        r.push_back(c.rgb.bl);

        depthObs.push_back(d);
        rgbObs.push_back(r);
    }

    int cameraFlags = useEstimate ? CV_CALIB_USE_INTRINSIC_GUESS : 0;

    // calibrate cameras individually
    vector<Mat> temp0,temp1;
    double error;
    printf("first camera\n");
    error = calibrateCamera( objectPoints,depthObs,Size(imgWidth,imgHeight),calib.firstK,calib.firstCoef,
                     temp0,temp1,cameraFlags);
    printf("  error = %f\n",error);
    printf("second\n");
    error = calibrateCamera( objectPoints,rgbObs,Size(imgWidth,imgHeight),calib.secondK,calib.secondCoef,
                     temp0,temp1,cameraFlags);
    printf("  error = %f\n",error);
    // calibrate stereo
    double reprojectionError = stereoCalibrate(objectPoints,depthObs,rgbObs,
                                               calib.firstK,calib.firstCoef,
                                               calib.secondK,calib.secondCoef,
                                               Size(imgWidth,imgHeight),
                                               calib.R,calib.T,calib.E,calib.F,
                                               TermCriteria(TermCriteria::COUNT+TermCriteria::EPS, 30, 1e-6),
                                               CV_CALIB_FIX_INTRINSIC);

    //TermCriteria(TermCriteria::COUNT+TermCriteria::EPS, 50, 1e-6),

    printf("* Calib T = %f  %f  %f\n",calib.T.at<double>(0),calib.T.at<double>(1),calib.T.at<double>(2));
    printf("* Coef2 %f  %f  %f\n",calib.secondCoef.at<double>(0),
           calib.secondCoef.at<double>(1),calib.secondCoef.at<double>(2));

    return reprojectionError;
}

int ComputeCalibrationRect::
        countInliers( const std::vector<CalibrationLocation>& obs , const StereoCalib &calib )
{
    int numInliers = 0;

    for( size_t i = 0; i < obs.size(); ++i ) {
        double error = computeError(obs.at(i),calib);

        if( error < inlierTol )
            numInliers++;
    }

    return numInliers;
}

void ComputeCalibrationRect::
        findInliers( const std::vector<CalibrationLocation>& obs , const StereoCalib &calib ,
                     std::vector<CalibrationLocation>& inliers )
{

    for( size_t i = 0; i < obs.size(); ++i ) {
        double error = computeError(obs.at(i),calib);

        if( error < inlierTol )
            inliers.push_back(obs.at(i));
    }
}

void ComputeCalibrationRect::
        pruneWorst( const std::vector<CalibrationLocation>& obs , const StereoCalib &calib ,
                    std::vector<CalibrationLocation>& best , double frac )
{
    std::vector<double> errors;
    for( size_t i = 0; i < obs.size(); ++i ) {
        errors.push_back(computeError(obs.at(i),calib));
    }

    sort (errors.begin(), errors.end());

    double threshold = errors.at( (int)(errors.size()*frac));

    for( size_t i = 0; i < obs.size(); ++i ) {
        double e = computeError(obs.at(i),calib);

        if( e < threshold ) {
            best.push_back(obs.at(i));
        }
    }
}

double ComputeCalibrationRect::
        computeError( const CalibrationLocation& obs , const StereoCalib &calib )
{
    // create vectors containing target corner pixel locations in each camera
    vector<cv::Point2f> cornersPixel1;
    cornersPixel1.push_back(obs.depth.tl);
    cornersPixel1.push_back(obs.depth.tr);
    cornersPixel1.push_back(obs.depth.br);
    cornersPixel1.push_back(obs.depth.bl);

    vector<cv::Point2f> cornersPixel2;
    cornersPixel2.push_back(obs.rgb.tl);
    cornersPixel2.push_back(obs.rgb.tr);
    cornersPixel2.push_back(obs.rgb.br);
    cornersPixel2.push_back(obs.rgb.bl);

    // compute target's location in the first camera's POV
    cv::Mat rvec;
    cv::Mat tvec;

    cv::solvePnP(cv::Mat(targetDesc),cv::Mat(cornersPixel1),calib.firstK,calib.firstCoef,rvec,tvec);

    cv::Mat rotMat(3,3,CV_64F);
    Rodrigues(rvec,rotMat);

    double maxError1 = reprojectionError(cornersPixel1,rotMat,tvec,calib.firstK,calib.firstCoef);

    // compute transform to second image
    rotMat = calib.R*rotMat;
    tvec = calib.R*tvec + calib.T;

    double maxError2 = reprojectionError(cornersPixel2,rotMat,tvec,calib.secondK,calib.secondCoef);

    printf("    Max error1 %f error2 %f\n",maxError1,maxError2);

    return maxError1 > maxError2 ? maxError1 : maxError2;
}

double ComputeCalibrationRect::reprojectionError( const std::vector<cv::Point2f> &pixels ,
                                                  const cv::Mat& R , const cv::Mat& T ,
                                                  const cv::Mat& K, const cv::Mat& coef ) {
    double maxError = 0;

    std::vector<cv::Point2f> projPoints;
    cv::projectPoints(cv::Mat(targetDesc),R,T,K,coef,projPoints);

    for( size_t i = 0; i < pixels.size(); ++i ) {
        const cv::Point2f &o = pixels.at(i);
        cv::Point2f &e = projPoints.at(i);

        double error = idch::distance<float>(o.x,o.y,e.x,e.y);

        if( error > maxError )
            maxError = error;
    }

    return maxError;
}

